from __future__ import absolute_import
from . import _nne_ops


#################################################################
# input parameter # type          # detail                      #
#################################################################
# input_data      # np.ndarray    # np.int8  | "NCHW"           #
# kernel_shape    # list          # [H, W]                      #
# strides         # list          # [H, W]                      #
# padding_size    # list          # [top, bottom, left, right]  #
# scale           # float         #
#################################################################
# output parameter                                              #
#################################################################
# output_data     # np.ndarray    # np.int8  | "NCHW"           #
#################################################################
def avepool(input_data, kernel_shape, strides, padding_size, scale):
    return _nne_ops._pool("ave", input_data, kernel_shape, strides, padding_size, scale)


#################################################################
# input parameter # type          # detail                      #
#################################################################
# input_data      # np.ndarray    # np.int8  | "NCHW"           #
# kernel_shape    # list          # [H, W]                      #
# strides         # list          # [H, W]                      #
# padding_size    # list          # [top, bottom, left, right]  #
#################################################################
# output parameter                                              #
#################################################################
# output_data     # np.ndarray    # np.int8  | "NCHW"           #
#################################################################
def maxpool(input_data, kernel_shape, strides, padding_size):
    return _nne_ops._pool("max", input_data, kernel_shape, strides, padding_size)


#################################################################
# input parameter # type          # detail                      #
#################################################################
# input_data      # np.ndarray    # np.int8  | "NCHW"           #
# weight_data     # np.ndarray    # np.int8  | "OIHW"           #
# bias_data  # np.ndarray or None # np.int32 | "O"              #
# scale_data      # np.ndarray    # np.float | "O"              #
# strides         # list          # [H, W]                      #
# padding_size    # list          # [top, bottom, left, right]  #
# activation_type # str           # 'relu' / 'lrelu' / 'none'   #
# lrelu_lmbda     # float         # optional                    #
# out_shift       # int in[-8,8]  # optional                    #
#################################################################
# output parameter                                              #
#################################################################
# output_data     # np.ndarray    # np.int8  | "NCHW"           #
#################################################################
def conv2d(input_data, weight_data, bias_data, scale_data, strides, padding_size, activation_type="relu",
           lrelu_lmbda=None, out_shift=0):
    return _nne_ops._conv("conv2d", input_data, weight_data, bias_data, scale_data, strides, padding_size,
                          activation_type, lrelu_lmbda, out_shift)


def dw_conv2d(input_data, weight_data, bias_data, scale_data, strides, padding_size, activation_type="relu",
              lrelu_lmbda=None, out_shift=0):
    return _nne_ops._conv("dw_conv2d", input_data, weight_data, bias_data, scale_data, strides, padding_size,
                          activation_type, lrelu_lmbda, out_shift)


#################################################################
# input parameter # type          # detail                      #
#################################################################
# input_data      # np.ndarray    # np.int8  | "NCHW"           #
# scale           # float         # (0,256)                     #
# activation_type # str           # 'relu' / 'lrelu'            #
# lrelu_lmbda     # float         # optional                    #
#################################################################
# output parameter                                              #
#################################################################
# output_data     # np.ndarray    # np.int8  | "NCHW"           #
#################################################################
def activation(input_data, scale, activation_type="relu", lrelu_lmbda=None):
    return _nne_ops._activation(input_data, scale, activation_type, lrelu_lmbda)


#################################################################
# input parameter # type          # detail                      #
#################################################################
# input_data      # np.ndarray    # np.int8  | "NI"             #
# weight_data     # np.ndarray    # np.int8  | "OI"             #
# bias_data  # np.ndarray or None # np.int32 | "O"              #
# scale_data      # np.ndarray    # np.float | "O"              #
# activation_type # str           # 'relu' / 'lrelu' / 'none'   #
# lrelu_lmbda     # float         # optional                    #
# out_shift       # int in[-8,8]  # optional                    #
#################################################################
# output parameter                                              #
#################################################################
# output_data     # np.ndarray    # np.int8  | "NO"             #
#################################################################
def fc(input_data, weight_data, bias_data, scale_data, activation_type="relu", lrelu_lmbda=None, out_shift=0):
    return _nne_ops._fc(input_data, weight_data, bias_data, scale_data, activation_type, lrelu_lmbda, out_shift)


#################################################################
# input parameter # type          # detail                      #
#################################################################
# input_1_data    # np.ndarray    # np.int8  | "NCHW"           #
# input_1_data    # np.ndarray    # np.int8  | "NCHW"           #
# scale_1         # float         #
# scale_2         # float         #
#################################################################
# output parameter                                              #
#################################################################
# output_data     # np.ndarray    # np.int8  | "NCHW"           #
#################################################################
def ew_add(input_1_data, input_2_data, scale_1, scale_2):
    return _nne_ops._ew("add", input_1_data, input_2_data, scale_1, scale_2)


#################################################################
# input parameter # type          # detail                      #
#################################################################
# input_1_data    # np.ndarray    # np.int8  | "NCHW"           #
# input_1_data    # np.ndarray    # np.int8  | "NCHW"           #
# scale           # float         #
#################################################################
# output parameter                                              #
#################################################################
# output_data     # np.ndarray    # np.int8  | "NCHW"           #
#################################################################
def ew_mul(input_1_data, input_2_data, scale):
    return _nne_ops._ew("mul", input_1_data, input_2_data, scale)


#################################################################
# input parameter # type          # detail                      #
#################################################################
# input_data      # np.ndarray    # np.int8  | "NCHW"           #
# scale           # float         #
#################################################################
# output parameter                                              #
#################################################################
# output_data     # np.ndarray    # np.int8  | "NCHW"           #
#################################################################
def global_average_pooling(input_data, scale):
    return _nne_ops._global_average_pooling(input_data, scale)


#################################################################
# input parameter # type          # detail                      #
#################################################################
# input_data      # np.ndarray    # np.int8  | "NCHW"           #
# h_factor        # int                                         #
# w_factor        # int                                         #
#################################################################
# output parameter                                              #
#################################################################
# output_data     # np.ndarray    # np.int8  | "NCHW"           #
#################################################################
def upsampling(input_data, h_factor, w_factor):
    return _nne_ops._upsampling(input_data, h_factor, w_factor)
